#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
标签信息表模型 (LabelInfo)
"""

from sqlalchemy import Column, String, Integer, DateTime, Index
from sqlalchemy.sql import func
from typing import List, Optional

from .BaseModel import BaseModel, standardize_date_format

class LabelInfo(BaseModel):
    """股票标签关联表模型"""
    __tablename__ = 'labelinfo'
    
    # 列名映射（CSV列名 -> 数据库字段名）
    aliases = {
        'date': 'date',
        'securityId': 'security_id',
        'label': 'label',
        'labelType': 'label_type',
        'confidence': 'confidence',
        '日期': 'date',
        '证券代码': 'security_id',
        '标签': 'label',
        '标签类型': 'label_type',
        '置信度': 'confidence'
    }
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    date = Column(String(10), nullable=False, comment='日期')
    security_id = Column(String(20), nullable=False, comment='证券代码')
    label = Column(String(100), nullable=False, comment='标签名称')
    label_type = Column(String(50), comment='标签类型')
    confidence = Column(Integer, default=100, comment='置信度：0-100')
    created_at = Column(DateTime, default=func.current_timestamp())
    
    # 索引
    __table_args__ = (
        Index('idx_label_info_date', 'date'),
        Index('idx_label_info_security', 'security_id'),
        Index('idx_label_info_label', 'label'),
        Index('idx_label_info_date_security', 'date', 'security_id'),
    )
    
    @classmethod
    def get_session(cls):
        """获取数据库会话"""
        return cls._db_manager.get_stock_data_session()
    
    @classmethod
    def _get_existing_record(cls, session, data):
        """获取现有记录"""
        return session.query(cls).filter(
            cls.date == data['date'],
            cls.security_id == data['security_id'],
            cls.label == data['label']
        ).first()
    
    @classmethod
    def _update_existing_record(cls, existing, data):
        """更新现有记录"""
        existing.label_type = data.get('label_type', existing.label_type)
        existing.confidence = data.get('confidence', existing.confidence)
    
    @classmethod
    def _to_dict(cls, result):
        """将结果转换为字典"""
        return {
            'id': result.id,
            'date': result.date,
            'security_id': result.security_id,
            'label': result.label,
            'label_type': result.label_type,
            'confidence': result.confidence,
            'created_at': result.created_at
        }
    
    @classmethod
    def get_labels_by_date(cls, date: str) -> List['LabelInfo']:
        """根据日期获取标签信息"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            results = session.query(cls).filter(cls.date == date).all()
            session.close()
            return results
        except Exception as e:
            # 获取日期标签失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def get_labels_by_security(cls, security_id: str, date: str = None) -> List['LabelInfo']:
        """根据证券代码获取标签信息"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            query = session.query(cls).filter(cls.security_id == security_id)
            if date:
                query = query.filter(cls.date == date)
            results = query.all()
            session.close()
            return results
        except Exception as e:
            # 获取证券标签失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def import_from_dataframe(cls, df):
        """从DataFrame导入标签信息数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            success_count = 0
            
            for _, row in df.iterrows():
                try:
                    # 标准化日期格式
                    date = standardize_date_format(row['date'])
                    if not date:
                        continue
                    
                    # 检查是否已存在相同记录
                    existing = session.query(cls).filter(
                        cls.date == date,
                        cls.security_id == row['security_id'],
                        cls.label == row['label']
                    ).first()
                    
                    if existing:
                        # 更新现有记录
                        cls._update_existing_record(existing, row.to_dict())
                    else:
                        # 创建新记录
                        label_info = cls(
                            date=date,
                            security_id=row['security_id'],
                            label=row['label'],
                            label_type=row.get('label_type', ''),
                            confidence=row.get('confidence', 100)
                        )
                        session.add(label_info)
                    
                    success_count += 1
                    
                except Exception as e:
                    # 处理标签信息数据失败: {str(e)}
                    continue
            
            session.commit()
            session.close()
            return success_count
            
        except Exception as e:
            # 导入标签信息数据失败: {str(e)}
            session.rollback()
            session.close()
            return 0
